import typer
from pathlib import Path
from utils import (
    run_extract_audio,
    run_whisper,
    extract_subtitle,
    translate_subtitle,
    embed_subtitle
)
import subprocess
import cv2
import numpy as np
from tqdm import tqdm

app = typer.Typer()

@app.command(name="extract-audio")
def extract_audio_cmd(input: Path):
    """提取音频为 audio.wav"""
    run_extract_audio(str(input))

@app.command(name="whisper")
def whisper_cmd(audio_path: Path = Path("audio.wav")):
    """使用 Whisper 识别音频为字幕"""
    run_whisper(str(audio_path))

@app.command(name="extract-sub")
def extract_sub_cmd(input: Path):
    """提取视频中的内嵌字幕"""
    extract_subtitle(str(input))

@app.command(name="translate")
def translate_cmd(srt_file: Path = Path("audio.srt")):
    """翻译字幕文件为中文"""
    translate_subtitle(str(srt_file))

@app.command(name="embed")
def embed_cmd(
    input: Path = typer.Argument(..., help="原始视频路径"),
    srt: Path = typer.Argument(..., help="字幕文件路径"),
    output: Path = typer.Option("output.mp4", help="输出视频路径")
):
    """将字幕嵌入视频中"""
    embed_subtitle(str(input), str(srt), str(output))

@app.command(name="remove-subtitle")
def remove_subtitle_cmd(
    input: Path = typer.Argument(..., help="输入原始视频文件路径"),
    output: Path = typer.Option("no_subtitle.mp4", help="输出无字幕视频路径")
):
    """移除视频中的内嵌字幕流"""
    print("🧹 正在移除原有字幕流...")
    cmd = ["ffmpeg", "-y", "-i", str(input), "-map", "0", "-map", "-0:s", "-c", "copy", str(output)]
    subprocess.run(cmd, check=True)

@app.command(name="inpaint-subtitle")
def inpaint_subtitle_cmd(
    input: Path = typer.Argument(..., help="带有烧录字幕的视频"),
    output: Path = typer.Option("cleaned.mp4", help="输出视频路径")
):
    """使用 OpenCV 对烧录字幕进行图像修复去除"""
    cap = cv2.VideoCapture(str(input))
    if not cap.isOpened():
        raise IOError("无法打开视频文件")

    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(str(output), fourcc, fps, (width, height))

    for _ in tqdm(range(total_frames), desc="处理视频帧"):
        ret, frame = cap.read()
        if not ret:
            break

        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
        mask_white = cv2.inRange(hsv, (0, 0, 180), (180, 50, 255))
        mask_yellow = cv2.inRange(hsv, (20, 100, 100), (35, 255, 255))
        mask = cv2.bitwise_or(mask_white, mask_yellow)
        mask[:int(height * 2 / 3), :] = 0

        inpainted = cv2.inpaint(frame, mask, 3, cv2.INPAINT_TELEA)
        out.write(inpainted)

    cap.release()
    out.release()
    print(f"✅ 烧录字幕清除完成：{output}")

@app.command(name="zit-remove-subtitle")
def zit_remove_cmd(
    input: Path = typer.Argument(..., help="带字幕的视频"),
    srt: Path = typer.Option(..., help="对应字幕文件"),
    output: Path = typer.Option("zit_cleaned.mp4", help="输出无字幕视频")
):
    """使用 ZITS 模型移除烧录字幕"""
    from utils import run_zits_removal
    run_zits_removal(str(input), str(srt), str(output))

@app.command()
def all(
    input: Path = typer.Option(..., help="输入视频文件"),
    output: Path = typer.Option("output.mp4", help="输出文件")
):
    """一键执行全部流程"""
    run_extract_audio(str(input))
    run_whisper("audio.wav")
    translate_subtitle("audio.srt")
    embed_subtitle(str(input), "translated.srt", str(output))

if __name__ == "__main__":
    app()
